import asyncio

from dotenv import load_dotenv
import os

from langchain.chains.question_answering.map_reduce_prompt import messages
from langchain.globals import set_verbose
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, BaseMessage
from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import AzureChatOpenAI
from langchain_community.chat_message_histories import ChatMessageHistory
from sqlalchemy.dialects.postgresql import JSONB
# from pprint import pprint
from rich import print as rprint
from rich.pretty import pprint

# from langchain_openai import AzureChatOpenAI
load_dotenv()

llm_model = AzureChatOpenAI(
    # openai_api_key=
    # openai_api_base=os.getenv("AZURE_OPENAI_ENDPOINT"),
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    azure_deployment=os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME"),
    api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
    temperature=0.7
)

# 初始化消息历史
history = ChatMessageHistory()


# 添加系统消息（可选）
# history.add_message(SystemMessage(content="将以下内容从英语翻译成中文"))

def main():
    msg = [
        SystemMessage(content="将下面内容从英文翻译成中文"),
        HumanMessage(content="It's a nice today")
    ]
    history.add_messages(msg)
    set_verbose(True)
    print("欢迎使用Azure OpenAI聊天助手！输入'退出'结束对话。")
    while True:
        user_input = input("你: ")

        if user_input.lower() == '退出':
            print("再见！")
            break

        # 添加用户消息到历史
        history.add_message(HumanMessage(content=user_input))
        parser = StrOutputParser()

        # 获取AI回复
        response = llm_model.invoke(history.messages)
        #
        print(f"原始数据: {response}")
        # 添加AI回复到历史
        history.add_message(AIMessage(content=response.content))
        print(f"AI助手: {response.content}")

        # parser
        result = parser.invoke(response)
        print(f"格式解析的结果：{result}")


"""
输出解析器
"""
# def out_parser():


"""
 链式调用
"""


def chain_call_demo() -> None:
    print("链式调用：")
    history.add_message(SystemMessage(content="将下面内容从英文翻译成中文"))
    while True:
        user_input = input("你: ")
        if user_input.lower() == '退出':
            print("再见！")
            break

        # 添加用户消息到历史
        history.add_message(HumanMessage(content=user_input))
        parser = StrOutputParser()
        # 获取AI回复
        chain = llm_model | parser
        response = chain.invoke(history.messages)

        # 添加AI回复到历史
        history.add_message(AIMessage(content=response))
        print(f"AI助手: {response}")


def stream_call_demo() -> None:
    """
     Stream流式输出调用
     ：同步调用，相当于open ai 参数里面的stream传递成true
    """
    print("====Stream流式调用======")
    chunks = []
    for chunk in llm_model.stream("天空是什么颜色？"):
        chunks.append(chunk)
        print(chunk.content, end="|", flush=True)


# astream
async def astream_call_demo() -> None:
    """
    astream 异步流式 调用
    """
    print("====AStream流式调用======")
    chunks = []
    async for chunk in llm_model.astream("天空是什么颜色？"):
        chunks.append(chunk)
        # 判断chunks长度为1的时候，打印chunks[0]
        if len(chunks) == 2:
            print(chunks[1])
        print(chunk.content, end="|", flush=True)


#  chain astream
async def chain_async_call_demo() -> None:
    print("====chain 异步调用======")
    chat_prompt = ChatPromptTemplate.from_template("给我讲一个关于{topic}的笑话")
    parser = StrOutputParser()
    chain = chat_prompt | llm_model | parser
    # chunks = []
    async for chunk in chain.astream({"topic": "鹦鹉"}):
        print(chunk, end="|", flush=True)


#  链式异步流式,且输出json
async def chain_async_call_json_demo() -> None:
    print("====chain 异步流式,且输出json======")
    parser = JsonOutputParser()
    chain = llm_model | parser
    async for text in chain.astream(
            "以JSON格式输出法国、西班牙和日本的国家及人口列表。使用一个待用countries 外部键的字典"):
        print(text, flush=True)


## stream event 事件流
async def stream_even_demo() -> None:
    print("====Stream事件流======")
    events = []
    async for event in llm_model.astream_events("hello", version="v2"):
        events.append(event)
    pprint(events)
    # pprint(events[:3])  ## 左闭右开，表示0，1，2


if __name__ == "__main__":
    # 设置日志级别 (DEBUG 查看更多细节)
    # logger.setLevel(logging.WARN)
    # 启动应用
    # main()

    # chain_call_demo()
    stream_call_demo()


    # asyncio.run(astream_call_demo())
    # asyncio.run(chain_async_call_json_demo())

    # asyncio.run(stream_even_demo())
